import tensorflow as tf
import cv2
import numpy as np

def conv_layer(x,filters,kernel_size,name):
    with tf.name_scope(name):
         conv = tf.layers.conv2d(x,filters=filters,kernel_size=kernel_size,strides=[1,1],padding='SAME',kernel_initializer=tf.contrib.layers.variance_scaling_initializer(), name=name+'_conv')
         conv= tf.nn.relu(conv,name=name+'_relu')
    return conv

def conv_bn_layer(x,filters,kernel_size,training,name):
    with tf.name_scope(name):
         conv = tf.layers.conv2d(x,filters=filters,kernel_size=kernel_size,strides=[1,1],padding='SAME',kernel_initializer=tf.contrib.layers.variance_scaling_initializer(), name=name+'_conv')
         conv = tf.layers.batch_normalization(conv,training=training,name=name+'_bn')
         conv= tf.nn.relu(conv,name=name+'_relu')
    return conv

def unpool(value,ind,name,ksize=[1,2,2,1]):
    with tf.variable_scope(name):
         #input_shape = pool.get_shape().as_list()
         #output_shape = (input_shape[0], input_shape[1] * ksize[1], input_shape[2] * ksize[2], input_shape[3])
         #flat_input_size = np.prod(input_shape)
         #flat_output_shape = [output_shape[0], output_shape[1] * output_shape[2] * output_shape[3]]
         #pool_ = tf.reshape(pool, [flat_input_size])
         #batch_range = tf.reshape(tf.range(output_shape[0], dtype=ind.dtype), shape=[input_shape[0], 1, 1, 1])
         #b = tf.ones_like(ind) * batch_range
         #b = tf.reshape(b, [flat_input_size, 1])
         #ind_ = tf.reshape(ind, [flat_input_size, 1])
         #ind_ = tf.concat([b, ind_], 1)
         #ret = tf.scatter_nd(ind_, pool_, shape=flat_output_shape)
         #ret = tf.reshape(ret, output_shape)

        sh = value.get_shape().as_list()
        dim = len(sh[1:-1])
        out = (tf.reshape(value, [-1] + sh[-dim:]))
        for i in range(dim, 0, -1):
            out = tf.concat([out, tf.zeros_like(out)],i)
        #out_size = [-1] + [s * 2 for s in sh[1:-1]] + [sh[-1]]
        out_size = [-1] + [sh[i] * ksize[i] for i in range(1,3)] + [sh[-1]] 
        out = tf.reshape(out, out_size, name=name)

         
    return out

def matsolver(images,trimaps,name = 'matsolver',training=True):
    '''
    Args:
         images are [batch_size,h,w,3]
         trimaps are [batch_size,h,w,1]    
    '''
    with tf.variable_scope(name,reuse = tf.AUTO_REUSE):
         pool_parameters=[]
         all_inputs = tf.concat([images,trimaps],3)
         conv1_1 = conv_layer(all_inputs,64,[3,3],'conv1_1')
         conv1_2 = conv_layer(conv1_1,64,[3,3],'conv1_2')
         pool1,arg1 = tf.nn.max_pool_with_argmax(conv1_2,ksize=[1,2,2,1],strides=[1,2,2,1],padding='SAME',name='pool1')
         pool_parameters.append(arg1)

         conv2_1= conv_layer(pool1,128,[3,3],'conv2_1')
         conv2_2= conv_layer(conv2_1,128,[3,3],'conv2_2')
         pool2,arg2 = tf.nn.max_pool_with_argmax(conv2_2,ksize=[1,2,2,1],strides=[1,2,2,1],padding='SAME',name='pool2')
         pool_parameters.append(arg2)
         
         conv3_1= conv_layer(pool2,256,[3,3],'conv3_1')
         conv3_2= conv_layer(conv3_1,256,[3,3],'conv3_2')
         conv3_3= conv_layer(conv3_2,256,[3,3],'conv3_3')
         pool3,arg3 = tf.nn.max_pool_with_argmax(conv3_3,ksize=[1,2,2,1],strides=[1,2,2,1],padding='SAME',name='pool3')
         pool_parameters.append(arg3)



         conv4_1= conv_layer(pool3,512,[3,3],'conv4_1')
         conv4_2= conv_layer(conv4_1,512,[3,3],'conv4_2')
         conv4_3= conv_layer(conv4_2,512,[3,3],'conv4_3')
         pool4,arg4 = tf.nn.max_pool_with_argmax(conv4_3,ksize=[1,2,2,1],strides=[1,2,2,1],padding='SAME',name='pool4')
         pool_parameters.append(arg4)




         conv5_1= conv_layer(pool4,512,[3,3],'conv5_1')
         conv5_2= conv_layer(conv5_1,512,[3,3],'conv5_2')
         conv5_3= conv_layer(conv5_2,512,[3,3],'conv5_3')
         pool5,arg5 = tf.nn.max_pool_with_argmax(conv5_3,ksize=[1,2,2,1],strides=[1,2,2,1],padding='SAME',name='pool5')
         pool_parameters.append(arg5)

         conv6_1 = conv_layer(pool5,4096,[7,7],'conv6_1')
         deconv6 = conv_bn_layer(conv6_1,512,[1,1],training,'deconv6')
         
         deconv5_1 = unpool(deconv6,pool_parameters[-1],name='deconv5_1')
         deconv5_2 = conv_bn_layer(deconv5_1,512,[5,5],training,'deconv5_2')

         deconv4_1 = unpool(deconv5_2,pool_parameters[-2],name='deconv4_1')
         deconv4_2 = conv_bn_layer(deconv4_1,256,[5,5],training,'deconv4_2')

         deconv3_1 = unpool(deconv4_2,pool_parameters[-3],name='deconv3_1')
         deconv3_2 = conv_bn_layer(deconv3_1,128,[5,5],training,'deconv3_2')

         deconv2_1 = unpool(deconv3_2,pool_parameters[-4],name='deconv2_1')
         deconv2_2 = conv_bn_layer(deconv2_1,64,[5,5],training,'deconv2_2')

         deconv1_1 = unpool(deconv2_2,pool_parameters[-5],name='deconv1_1')
         deconv1_2 = conv_bn_layer(deconv1_1,64,[5,5],training,'deconv1_2')

         pred = conv_layer(deconv1_2,1,[5,5],'pred')
         pred_mattes = tf.nn.sigmoid(pred,name='predictions')

    return pred_mattes


if __name__=="__main__":
   image=cv2.imread('./data/image.png')
   trimap = cv2.imread('./data/image.png',0)
   h,w = trimap.shape
   trimap = np.reshape(trimap,[h,w,1])
   im_placeholder = tf.placeholder(tf.float32,[None,h,w,3])
   trimap_placeholder = tf.placeholder(tf.float32,[None,h,w,1])
   mattes = matsolver(im_placeholder,trimap_placeholder)
   sess=tf.Session()
   sess.run(tf.global_variables_initializer())
   predict_mattes = sess.run(mattes,feed_dict={im_placeholder:[image],trimap_placeholder:[trimap]})
   print(predict_mattes.shape)
